package 困难.转换思路;

import java.util.HashMap;
import java.util.Map;

/**
 * 给你一个长度为 n 的数组 nums ，该数组由从 1 到 n 的
 * 不同 整数组成。另给你一个正整数 k 。
 * 统计并返回 num 中的 中位数 等于 k 的非空子数组的数目。
 * 注意：
 * 数组的中位数是按 递增 顺序排列后位于 中间 的那个元素，如果
 * 数组长度为偶数，则中位数是位于中间靠 左 的那个元素。
 * 例如，[2,3,1,4] 的中位数是 2 ，[8,4,3,5,1] 的中位数是 4 。
 * 子数组是数组中的一个连续部分。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode.cn/problems/count-subarrays-with-median-k
 */
public class 统计中位数为K的子数组_6248 {

    public static void main(String[] args) {

        System.out.println(countSubarrays2(new int[]{3, 2, 1, 4, 5}, 4));

    }

    /**
     * 模拟
     * 1，数组中包含 1 到 n，可以先找到 k 所在的索引 kIdx，子数组必须包含 kIdx
     * 2，可以由 kIdx 向左向右扩张，并统计 > k 的个数 big 和 < k 的个数 less
     * 当 big==less 或者 big+1==less 时，中位数是 k
     */
    public static int countSubarrays(int[] nums, int k) {
        int kIdx = 0;
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == k) {
                kIdx = i;
                break;
            }
        }

        int count = 0;
        for (int i = kIdx; i < nums.length; i++) {
            int big = 0, less = 0;
            for (int j = i; j >= 0; j--) {
                if (nums[j] > k) {
                    big++;
                } else if (nums[j] < k) {
                    less++;
                }
                if (j <= kIdx && big - less <= 1 && big >= less) {
                    count++;
                }
            }
        }
        return count;
    }

    /**
     * 将 >k 的数变成 1 ，<k 的数变成 -1，=k 的数变成0
     * 可以将问题转换为，包含索引 kIdx 的连续子数组和
     * =0 或者 =1 的个数是多少
     * 具体思路：
     * 用哈希表存储，i < kIdx ,prefixSum[i] 表示 [i,kIdx] 的和
     * 过程中每个 和 出现的个数
     * 然后遍历 [kIdx,nums.length-1] ,维护 suffixSum ，表示 [kIdx,j] j<=nums.length-1
     * 遍历过程中的和
     * 如果哈希表中存在 -suffixSum 的前半部分和，使得加上后半部分的和suffixSum = 0
     * 即 prefixSum+suffixSum = 0
     * 或者 prefixSum+suffixSum = 1
     * 求的是哈希表中 map[-suffixSum]+map[1-suffixSum] 的数量
     */
    public static int countSubarrays2(int[] nums, int k) {
        int kIdx = 0;
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] == k) {
                kIdx = i;
            }
        }
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] > k) {
                nums[i] = 1;
            } else if (nums[i] == k) {
                nums[i] = 0;
            } else {
                nums[i] = -1;
            }
        }
        Map<Integer, Integer> prefixSumMap = new HashMap<>();
        int prefixSum = 0;
        for (int i = kIdx; i >= 0; i--) {
            prefixSum += nums[i];
            prefixSumMap.put(prefixSum, prefixSumMap.getOrDefault(prefixSum, 0) + 1);
        }
        int suffixSum = 0;
        int count = 0;
        for (int i = kIdx; i < nums.length; i++) {
            suffixSum += nums[i];
            count += prefixSumMap.getOrDefault(-suffixSum, 0);
            count += prefixSumMap.getOrDefault(1 - suffixSum, 0);
        }

        return count;
    }

}
